
#include "stdafx.h"
#include "EnumGdi.h"
#include "LogClass.h"
#include "tick.h"

GDIType GetGDIType(USHORT type)
{
	switch (type)
	{
	case 1:		return GDI_DC;
	case 4:		return GDI_Region;
	case 5:		return GDI_Bitmap;
	case 8:		return GDI_Palette;
	case 10:	return GDI_Font;
	case 16:	return GDI_Brush;
	case 48:	return GDI_Pen;
	}
	return GDI_Unknown;
}

stGDIResult * GetGDIResult_0()
{
	HANDLE hProcess = OpenProcess(PROCESS_QUERY_INFORMATION | PROCESS_VM_READ, FALSE, GetCurrentProcessId());
	// determine if 64 or 32-bit processor
	SYSTEM_INFO si;
	GetNativeSystemInfo(&si);

	// NOTE: as this is undocumented, it *may vary* depending on bitness (32/64) and on Windows version.
	// use WinDbg "dt ntdll!_PEB" command and search for GdiSharedHandleTable offset to find the truth out
	size_t GdiSharedHandleTableOffset = si.wProcessorArchitecture == PROCESSOR_ARCHITECTURE_AMD64 ? 0xF8 : 0x94;
	size_t tableCount = 16384; // count of GDI table cells

	// determine if this process is running on WOW64
	BOOL wow;
	IsWow64Process(hProcess, &wow);

	// read basic info to get PEB address, we only need the beginning of PEB
	size_t pebSize = GdiSharedHandleTableOffset + 8;
	LPBYTE peb = (LPBYTE)HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, pebSize);
	LPVOID pTable = nullptr;
	stGDIResult *result = nullptr;
	do {
		if (wow)
		{
			// we're running as a 32-bit process in a 64-bit process
			PROCESS_BASIC_INFORMATION_WOW64 pbi;
			ZeroMemory(&pbi, sizeof(pbi));

			// get process information from 64-bit world
			_NtQueryInformationProcess query = (_NtQueryInformationProcess)GetProcAddress(GetModuleHandleA("ntdll.dll"), "NtWow64QueryInformationProcess64");
			if (query(hProcess, 0, &pbi, sizeof(pbi), NULL) != 0) break;

			// read PEB from 64-bit address space
			_NtWow64ReadVirtualMemory64 read = (_NtWow64ReadVirtualMemory64)GetProcAddress(GetModuleHandleA("ntdll.dll"), "NtWow64ReadVirtualMemory64");
			if (read(hProcess, pbi.PebBaseAddress, peb, pebSize, NULL) != 0) break;

			// get GDI table ptr from PEB
			GDICELL_WOW64* gdiTable = (GDICELL_WOW64*)*(LPVOID*)(peb + GdiSharedHandleTableOffset); // address in remote process adress space

			size_t tableSize = sizeof(GDICELL_WOW64) * tableCount; // size of GDI table
			pTable = HeapAlloc(GetProcessHeap(), 0, tableSize); if (!pTable) break;
			GDICELL_WOW64* table = (GDICELL_WOW64*)pTable; // local table copied over to our address space

			// copy GDI table
			if (read(hProcess, gdiTable, table, tableSize, NULL) != 0) break;
			result = (stGDIResult *)HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, sizeof(stGDIResult)); if (!result) break;
			for (size_t i = 0; i < tableCount; i++) if (table[i].wProcessId == GetCurrentProcessId())
			{
				GDIType type = GetGDIType(table[i].wType & 0x7F);
				result->ary[type].count++;
				result->totalCount++;
			}
			intptr_t pos[8];
			for (intptr_t i = 0; i < 8; i++)
			{
				if (result->ary[i].count)
				{
					result->ary[i].handles = (HGDIOBJ*)HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, result->ary[i].count*sizeof(HGDIOBJ));
					result->ary[i].same = (bool*)HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, result->ary[i].count*sizeof(bool));
				}
				pos[i] = 0;
			}
			for (size_t i = 0; i < tableCount; i++) if (table[i].wProcessId == GetCurrentProcessId())
			{
				HGDIOBJ gdiHandle = (HGDIOBJ)((table[i].wUpper << 16) + i);
				GDIType type = GetGDIType(table[i].wType & 0x7F);
				HGDIOBJ *buffer = result->ary[type].handles;
				if (pos[type] < result->ary[type].count && buffer)
				{
					intptr_t j = 0; for (; j < pos[type]; j++) if (gdiHandle < buffer[j]) break;
					for (intptr_t k = pos[type]; k > j; k--) buffer[k] = buffer[k - 1];
					buffer[j] = gdiHandle; pos[type]++;
				}
			}
		}
		else
		{
			// we're running as a 32-bit process in a 32-bit OS, or as a 64-bit process in a 64-bit OS
			PROCESS_BASIC_INFORMATION pbi;
			ZeroMemory(&pbi, sizeof(pbi));

			// get process information
			_NtQueryInformationProcess query = (_NtQueryInformationProcess)GetProcAddress(GetModuleHandleA("ntdll.dll"), "NtQueryInformationProcess");
			if (query(hProcess, 0, &pbi, sizeof(pbi), NULL) != 0) break;

			// read PEB
			_NtReadVirtualMemory read = (_NtReadVirtualMemory)GetProcAddress(GetModuleHandleA("ntdll.dll"), "NtReadVirtualMemory");
			if (read(hProcess, pbi.PebBaseAddress, peb, pebSize, NULL) != 0) break;

			// get GDI table ptr
			GDICELL* gdiTable = (GDICELL*)*(LPVOID*)(peb + GdiSharedHandleTableOffset); // address in remote process adress space
			if (gdiTable == NULL) break;

			size_t tableSize = sizeof(GDICELL) * tableCount; // size of GDI table
			pTable = HeapAlloc(GetProcessHeap(), 0, tableSize); if (!pTable) break;
			GDICELL* table = (GDICELL*)pTable; // local table copied over to our address space

			// read GDI table
			if (read(hProcess, gdiTable, table, tableSize, NULL) != 0) break;
			result = (stGDIResult *)HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, sizeof(stGDIResult)); if (!result) break;
			for (size_t i = 0; i < tableCount; i++) if (table[i].wProcessId == GetCurrentProcessId())
			{
				GDIType type = GetGDIType(table[i].wType & 0x7F);
				result->ary[type].count++;
				result->totalCount++;
			}
			intptr_t pos[8];
			for (intptr_t i = 0; i < 8; i++)
			{
				if (result->ary[i].count)
				{
					result->ary[i].handles = (HGDIOBJ*)HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, result->ary[i].count*sizeof(HGDIOBJ));
					result->ary[i].same = (bool*)HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, result->ary[i].count*sizeof(bool));
				}
				pos[i] = 0;
			}
			for (size_t i = 0; i < tableCount; i++) if (table[i].wProcessId == GetCurrentProcessId())
			{
				HGDIOBJ gdiHandle = (HGDIOBJ)((table[i].wUpper << 16) + i);
				GDIType type = GetGDIType(table[i].wType & 0x7F);
				HGDIOBJ *buffer = result->ary[type].handles;
				if (pos[type] < result->ary[type].count && buffer)
				{
					intptr_t j = 0; for (; j < pos[type]; j++) if (gdiHandle < buffer[j]) break;
					for (intptr_t k = pos[type]; k > j; k--) buffer[k] = buffer[k - 1];
					buffer[j] = gdiHandle; pos[type]++;
				}
			}
		}
	} while (0);
	CloseHandle(hProcess);
	if (peb) HeapFree(GetProcessHeap(), 0, peb);
	if (pTable) HeapFree(GetProcessHeap(), 0, pTable);
	return result;
}

LPCWSTR typeCaption(GDIType type)
{
	switch (type)
	{
	case GDI_DC:		return L"DC";
	case GDI_Region:	return L"Region";
	case GDI_Bitmap:	return L"Bitmap";
	case GDI_Palette:	return L"Palette";
	case GDI_Font:		return L"Font";
	case GDI_Brush:		return L"Brush";
	case GDI_Pen:		return L"Pen";
	}
	return L"Unknown";
}
void PrintOut(LPCWSTR caption, const stGDIArray *ary)
{
	if (!ary->handles)
		gLog.add(L"%s is null", caption);
	else
	{
		gLog.add(L"%s Count=%d", caption, ary->count);
		for (intptr_t j = 0; j < ary->count;)
		{
			wchar_t v[1024]; v[0] = 0;
			for (intptr_t k = 0; k < 8 && j < ary->count; k++, j++)
			{
				if (k) wcscat_s(v, L" ");
				wchar_t v0[20];
#ifdef _WIN64
				swprintf_s(v0, L"%s%16I64x", ary->same[j] ? L" " : L"*", (intptr_t)ary->handles[j]);
#else
				swprintf_s(v0, L"%s%08x", ary->same[j] ? L" " : L"*", (int)ary->handles[j]);
#endif
				wcscat_s(v, v0);
			}
			gLog.add(L"  %s", v);
		}
	}
}
void PrintOut(LPCWSTR caption, const stGDIResult *result)
{
	if (!result)
		gLog.add(L"%s is null", caption);
	else
	{
		gLog.add(L"%s TotalCount=%d", caption, result->totalCount);
		for (intptr_t i = 0; i < 8; i++) PrintOut(typeCaption((GDIType)i), &result->ary[i]);
	}
}
void FreeResult(stGDIResult *result)
{
	if (!result) return;
	for (intptr_t i = 0; i < 8; i++)
	{
		if (result->ary[i].handles) HeapFree(GetProcessHeap(), 0, result->ary[i].handles);
		if (result->ary[i].same) HeapFree(GetProcessHeap(), 0, result->ary[i].same);
	}
	HeapFree(GetProcessHeap(), 0, result);
}
void PrintOut(stGDIResult *before, stGDIResult *after, bool freeBefore, bool freeAfter)
{
	gLog.add(L"-----------------------------------");
	if (!before || !after)
	{
		PrintOut(L"Before", before);
		PrintOut(L"After", after);
	}
	else
	{
		bool isSame = true;
		for (intptr_t i = 0; i < 8; i++)
		{
			bool issm; stGDIArray *bAry = &before->ary[i], *aAry = &after->ary[i];
			if (bAry->count != aAry->count)
				issm = false;
			else
			{
				issm = true;
				if (i != GDI_Region && i != GDI_Brush) for (intptr_t j = 0; j < bAry->count; j++) if (bAry->handles[j] != aAry->handles[j])
				{
					issm = false; break;
				}
			}
			if (!issm)
			{
				for (intptr_t j = 0, k = 0; j < bAry->count || k < aAry->count;)
				{
					if (j >= bAry->count)
					{
						aAry->same[k] = false; k++;
					}
					else if (k >= aAry->count)
					{
						bAry->same[j] = false; j++;
					}
					else if (bAry->handles[j] == aAry->handles[k])
					{
						bAry->same[j] = aAry->same[k] = true; j++; k++;
					}
					else if (bAry->handles[j] < aAry->handles[k])
					{
						bAry->same[j] = false; j++;
					}
					else
					{
						aAry->same[k] = false; k++;
					}
				}
				gLog.add(L"%s isn't same(%d,%d)", typeCaption((GDIType)i), bAry->count, aAry->count);
				PrintOut(L"Before", bAry);
				PrintOut(L"After", aAry);
				isSame = false;
			}
		}
		if (isSame) gLog.add(L"Before and After is same");
	}
	gLog.flush();
	if (freeBefore && before) FreeResult(before);
	if (freeAfter && after) FreeResult(after);
}

bool isSameResult(const stGDIResult*a, const stGDIResult*b)
{
	if (a->totalCount != b->totalCount) return false;
	for (intptr_t i = 0; i < 8; i++)
	{
		if (a->ary[i].count != b->ary[i].count) return false;
		for (intptr_t j = 0; j < a->ary[i].count; j++) if (a->ary[i].handles[j] != b->ary[i].handles[j]) return false;
	}
	return true;
}

extern void IDLE();
stGDIResult * GetGDIResult()
{
	stGDIResult *n = GetGDIResult_0(); if (!n) return n;
	for (TICK tk = TICK::Now(5000);; IDLE())
	{
		stGDIResult *m = GetGDIResult_0(); if (!m) return n;
		if (isSameResult(m, n))
		{
			FreeResult(m); if (TICK::Now() >= tk) return n;
		}
		else
		{
			FreeResult(n); n = m; tk = TICK::Now(5000);
		}
	}
}
